import abc
import numpy as np


class LMModel(metaclass=abc.ABCMeta):
    def __init__(self, context_length: int):
        self.__context_length = context_length

    @property
    def output_logits(self) -> bool:
        raise NotImplementedError

    @property
    def context_length(self) -> int:
        return self.__context_length

    @abc.abstractmethod
    def predict(
            self, inputs: np.ndarray, batch_size: int = 1,
            **kwargs) -> np.ndarray:
        pass

    @staticmethod
    def format_inputs(inputs: np.ndarray):
        if len(inputs.shape) > 1:
            assert len(inputs.shape) == 2 and inputs.shape[0] == 1
        else:
            assert len(inputs.shape) == 1
            inputs = inputs.reshape((1, -1))
        return inputs


class GenerateLMModel(LMModel, metaclass=abc.ABCMeta):
    def __init__(self, context_length: int):
        super().__init__(context_length)

    @abc.abstractmethod
    def generate(self, inputs: np.ndarray, generate_count: int,
                 **kwargs) -> np.ndarray:
        pass
